import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d as p3d

'''
Rossler偏微分方程组为
dx=−(y+z)
dy=x+ay 
dz=b+z(x−c)
'''
# 设置Rossler常微分方程组的系统参数
a = 0.5
b = 2.0
c = 4.0
# dx=−(y+z)
def fx(x, y, z):
    return - y - z
# dy=x+ay
def fy(x, y, z):
    return x + a * y
# dz=b+z(x−c)
def fz(x, y, z):
    return b + z * (x - c)


# 定义函数Rossler求解方程组的解
def Rossler(x0, y0, z0, N):
    h = 0.01
    x = []
    y = []
    z = []
    # 四阶龙格库塔方法求解Lorenz常微分方程组的解，T为迭代轮次
    for i in range(N):
        K1 = fx(x0, y0, z0)
        L1 = fy(x0, y0, z0)
        M1 = fz(x0, y0, z0)

        K2 = fx(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)
        L2 = fy(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)
        M2 = fz(x0 + h * K1 / 2, y0 + h * L1 / 2, z0 + h * M1 / 2)

        K3 = fx(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)
        L3 = fy(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)
        M3 = fz(x0 + h * K2 / 2, y0 + h * L2 / 2, z0 + h * M2 / 2)

        K4 = fx(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        L4 = fy(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        M4 = fz(x0 + h * K3, y0 + h * L3, z0 + h * M3)
        # x0、y0、z0统一更新
        x0 = x0 + h / 6 * (K1 + 2 * K2 + 2 * K3 + K4)
        y0 = y0 + h / 6 * (L1 + 2 * L2 + 2 * L3 + L4)
        z0 = z0 + h / 6 * (M1 + 2 * M2 + 2 * M3 + M4)

        x.append(x0)
        y.append(y0)
        z.append(z0)
    return x, y, z


def main():
    T = 20000
    # 设初值
    x0, y0, z0 = 1, 0, 0
    x, y, z = Rossler(x0, y0, z0, T)
    ax = plt.subplot(121, projection="3d")
    ax.scatter(x, y, z, s=5)
    ax.set_xlabel('x(t)')
    ax.set_ylabel('y(t)')
    ax.set_zlabel('z(t)')
    ax.set_title('x0 = 1 y0 = 0 z0 = 0')

    ax.grid(False)
    # 设置微小变化的初值
    x0 = 1
    y0 = 0
    z0 = 0.00001
    xx, yy, zz = Rossler(x0, y0, z0, T)
    ax = plt.subplot(122, projection="3d")
    ax.scatter(xx, yy, zz, s=5)
    ax.set_xlabel('x(t)')
    ax.set_ylabel('y(t)')
    ax.set_zlabel('z(t)')
    ax.set_title('x0 = 1 y0 = 0 z0 = 0.0001')
    ax.grid(False)

    t = np.arange(0, T)
    plt.scatter(t, x, s=1)
    plt.scatter(t, xx, s=1, c='b')
    plt.title('x-xx')
    plt.xlabel('x(t)')
    plt.ylabel('y(t)')
    plt.savefig('Rossler的x(t)轨道演化.png')
    plt.show()

    plt.scatter(t, y, s=1)
    plt.scatter(t, yy, s=1, c='r')
    plt.title('y-yy')
    plt.xlabel('x(t)')
    plt.ylabel('y(t)')
    plt.savefig('Rossler的y(t)轨道演化.png')
    plt.show()

    plt.scatter(t, z, s=1)
    plt.scatter(t, zz, s=1, c='g')
    plt.title('z-zz')
    plt.xlabel('x(t)')
    plt.ylabel('y(t)')
    plt.savefig('Rossler的z(t)轨道演化.png')
    plt.show()


if __name__ == '__main__':
    main()